import cv2
import numpy as np
import torch
from torch import FloatTensor, tensor
from torch.nn import (
    AdaptiveAvgPool2d,
    Conv2d,
    # Dropout层，用于在训练过程中随机忽略部分神经元，防止过拟合
    Dropout,
    # 全连接层，用于将输入特征映射到指定维度的输出
    Linear,
    # （1）所有神经网络模块的基类，自定义神经网络时需要继承该类
    Module,
    # 线性整流函数（ReLU）激活层，用于给神经网络引入非线性特性
    ReLU,
    # （2）顺序容器，用于按顺序组合多个神经网络模块
    Sequential,
    Sigmoid,
)
from torchvision.models import resnet18

from backend.model_core.global_variables import GlobalVariables


class DualFundusClassifier(Module):
    def __init__(self, num_classes=8):
        super().__init__()
        # 模型是在 ImageNet-1K 数据集 上训练的。
        # 这是该模型的一个 官方发布的预训练版本（v1）。
        # ImageNet-1K 包含约 128 万张图片，分为 1000 个类别。
        # 使用这些权重意味着你可以直接将模型用于图像分类任务，或者进行迁移学习
        # 这是 ResNet 架构设计决定的。ResNet-18 的通道扩展过程如下：
        # 层	输出通道数
        # Conv1	64
        # Layer1（残差块）	64
        # Layer2（下采样）	128
        # Layer3（下采样）	256
        # Layer4（下采样）	512
        resnet = resnet18(weights="IMAGENET1K_V1")
        self.features = Sequential(
            # in:[B, 2, 512, 512]	out:[B, 64, 256, 256]
            # 根据卷积公式计算输出尺寸：H_out = (H_in + 2*padding - kernel_size) / stride + 1
            # 此处 (512 + 2*3 - 7) / 2 + 1 = 256.5，由于PyTorch中默认向下取整，实际输出应为256
            Conv2d(  # 二维卷积层，常用于处理图像数据，进行特征提取
                in_channels=2,  # 输入通道: 2
                out_channels=64,  # 输出通道: 64
                kernel_size=7,  # 卷积核大小: 7
                stride=2,  # 步长: 2
                padding=3,  # 填充: 3
                bias=False,  # 是否使用偏置项: False
            ),
            # in:[B, 64, 256, 256]	out:[B, 512, 32, 32]
            # 标准ResNet18结构：经过4个残差块组进行下采样
            # Layer1: 64→64通道保持, 尺寸不变(256×256)
            # Layer2: 64→128通道翻倍, 尺寸减半(128×128)
            # Layer3: 128→256通道翻倍, 尺寸减半(64×64)
            # Layer4: 256→512通道翻倍, 尺寸减半(32×32)
            # AdaptiveAvgPool2d(1, 1) # 输出尺寸: (B, 512, 1, 1)
            *list(resnet.children())[
                1:-1
            ],  # 移除最后的全连接层Linear(in_features=512, out_features=num_classes)
        )

        # 简化结构化数据分支
        # 32是深度学习中的常见中间维度（如CNN中的通道数），在类似任务中已被验证能有效平衡表达能力和计算成本
        self.struct_branch = Sequential(
            Linear(2, 32),  # 输入特征数: 2，输出特征数: 32
            ReLU(),  # 输入输出特征数不变: 32
            Dropout(0.2),  # 输入输出特征数不变: 32
        )

        # 调整分类头维度计算
        self.classifier = Sequential(
            Linear(512 + 32, 256),  # 输入特征数: 512 + 32，输出特征数: 256
            ReLU(),  # 输入输出特征数不变: 256
            Dropout(0.3),  # 输入输出特征数不变: 256
            Linear(256, num_classes),  # 输入特征数: 256，输出特征数: num_classes
            Sigmoid(),  # 输入输出特征数不变: num_classes
        )

    def forward(self, x, age, sex):
        features = self.features(x)

        # 展平特征图
        features = features.view(features.size(0), -1)

        structured = torch.cat([age, sex], dim=1)

        # 应用结构化数据分支
        structured = self.struct_branch(structured)

        combined = torch.cat([features, structured], dim=1)

        return self.classifier(combined)


class DualFundusModel:
    def __init__(self, model_path="models/best_model.pth"):
        """
        初始化模型加载器

        参数:
            model_path: 模型权重文件路径
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = DualFundusClassifier().to(self.device)
        # 加载训练好的模型
        self.model.load_state_dict(
            torch.load(model_path, map_location=self.device, weights_only=True)
        )
        # 设置为评估模式
        self.model.eval()
        # 定义图像尺寸配置
        self.image_size = (512, 512)

    def _preprocess_image(self, image):
        """
        单张图像预处理: 直方图均衡化 + 自适应对比度增强

        参数:
            image: 输入BGR格式的numpy数组图像
        返回:
            预处理后的单通道归一化图像（形状: (512, 512, 1)）
        """
        # 假设输入是numpy数组格式的BGR图像
        if image is None:
            raise ValueError("无效的图像输入")

        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)

        # 图像归一化
        normalized = enhanced / 255.0

        # 调整尺寸
        resized = cv2.resize(normalized, self.image_size)

        # 添加通道维度
        return resized[..., np.newaxis]

    def predict(self, left_image, right_image, age, sex):
        """
        预测方法（接受内存中的图像数据）

        参数:
            left_image: 左眼图像numpy数组（BGR格式）
            right_image: 右眼图像numpy数组（BGR格式）
            age: 年龄值
            sex: 性别值 (0表示女性, 1表示男性)

        返回:
            list: 八个类别的概率列表 [N, D, G, C, A, H, M, O]
        """
        # 图像预处理
        left = self._preprocess_image(left_image)
        right = self._preprocess_image(right_image)

        # 拼接双目图像
        image = np.concatenate([left, right], axis=-1)

        # 添加维度转置 NHWC -> NCHW
        image = np.transpose(image, (2, 0, 1))  # 转置为(C,H,W)

        # 转换为张量并添加批量维度，模型直接接收单个样本输入，需手动添加批量维度unsqueeze(0)
        image_tensor = FloatTensor(image).unsqueeze(0).to(self.device)
        age_tensor = tensor([age/ GlobalVariables.max_age]).float().unsqueeze(0).to(self.device)
        sex_tensor = tensor([sex]).float().unsqueeze(0).to(self.device)

        # 模型推理
        with torch.no_grad():
            outputs = self.model(image_tensor, age_tensor, sex_tensor)

        # 返回概率结果
        return outputs.cpu().numpy()[0].tolist()
